#if _MSC_VER >= 1600
#pragma execution_character_set("utf-8")
#endif

#include "DefectDetection.h"

DefectDetection::DefectDetection()
{

}
DefectDetection::DefectDetection(string modelPath, int modelType, bool useBoxInfo, bool usePointInfo, string saveDir)
{
	// 设置sam得一些参数
	this->cfg = new Configuration();
	if (modelType == 1) {
		this->cfg->EncoderModelPath = modelPath + "/vit_b/encoder/sam_vit_b_01ec64_encoder.onnx";
		this->cfg->DecoderModelPath = modelPath + "/vit_b/decoder/sam_vit_b_01ec64_decoder.onnx";
	}
	else {
		this->cfg->EncoderModelPath = modelPath + "/vit_l/encoder/sam_vit_l_0b3195_encoder.onnx";
		this->cfg->DecoderModelPath = modelPath + "/vit_l/decoder/sam_vit_l_0b3195_decoder.onnx";
	}
	this->cfg->UseBoxInfo = useBoxInfo;
	this->cfg->UsePointInfo = usePointInfo;
	this->cfg->SaveDir = saveDir;
	this->cfg->SaveName = "temp";

	//创建 提取结果保存的文件夹 用时间戳表示每次的提取的路径
	QString saveDirPath = QString::fromStdString(this->cfg->SaveDir) + "/" + QString::number(QDateTime::currentDateTime().toTime_t()); // 获取当前时间戳
	QDir dir(saveDirPath);
	if (!dir.exists()) {
		if (!dir.mkdir(saveDirPath)) {
			std::cout << dir.absolutePath().toStdString() << "created failed!";
			return;
		}
	}
	this->cfg->SaveDir = saveDirPath.toStdString();
	// 初始化环境 加载sam模型
	this->samOnnxRunner = new SAMOnnxRunner();
	this->samOnnxRunner->InitOrtEnv(*(this->cfg));
	
}
DefectDetection::~DefectDetection()
{
	delete this->samOnnxRunner;
	delete this->cfg;
}

bool DefectDetection::YoloDefectDetection(string modelPath, string inputFileName,string saveFileName, vector<std::string> classes_Names, std::vector<int> classes_Ids, vector<string> classes_filters, vector<int> ids_filters)
{
	DCSP_CORE* yoloDetector = new DCSP_CORE;
	std::string model_path = modelPath;
	//read_coco_yaml(yoloDetector);
	yoloDetector->classes = classes_Names;//{ "arrow", "jiaoyi_tuoluo", "leiji", "louyou", "tatong_xiushi" };
	yoloDetector->ids = classes_Ids;//{ 0, 1, 2, 3, 4 };
#ifdef USE_CUDA
	// GPU FP32 inference
	DCSP_INIT_PARAM params{ model_path, YOLO_ORIGIN_V8, {640, 640},  0.1, 0.5, true };
	// GPU FP16 inference
	// DCSP_INIT_PARAM params{ model_path, YOLO_ORIGIN_V8_HALF, {640, 640},  0.1, 0.5, true };
#else
	// CPU inference
	DCSP_INIT_PARAM params{ model_path, YOLO_ORIGIN_V8, {640, 640}, 0.1, 0.5, false };
#endif
	yoloDetector->CreateSession(params);
	vector<string> classes_filter = classes_filters;//{ "louyou" };
	vector<int> ids_filter = ids_filters;//{ 3 };

	std::unordered_map<int, std::string> id_to_class;  // idclassӳϵ
	for (int i = 0; i < classes_filter.size(); i++) {
		id_to_class[ids_filter[i]] = classes_filter[i];
	}

	cv::Mat resized_img;
	cv::Mat imageMat = cv::imread(inputFileName);
	if (imageMat.empty())
	{
		return false;
	}
	//cv::resize(imageMat, resized_img, cv::Size(1024, 1024), cv::INTER_LINEAR);
	std::vector<DCSP_RESULT> res;
	yoloDetector->RunSession(imageMat, res);

	std::vector<DCSP_RESULT> res_filtered;
	for (auto& re : res) {
		if (id_to_class.find(re.classId) != id_to_class.end() && re.confidence >= 0.6) {
			res_filtered.push_back(re);
		}
	}
	std::string saveFileNameCSV = saveFileName + "_ResultCSV.csv";
	Utils utilsObj;
	utilsObj.persistResult(res_filtered, id_to_class, saveFileNameCSV);
	GdalOpencvTensor gdalOpencv;
	gdalOpencv.save2shp(inputFileName, res_filtered, id_to_class, saveFileName);
	return true;
}

bool DefectDetection::SegModel(string modelPath, string inputFileName, string saveFileName, int iflag)
{
	GeologicalHazardSeg* fcn_resnet101 = NULL;
	if (iflag==1){
		fcn_resnet101 = new GeologicalHazardSeg(modelPath, 16, 256, 256);
	}
	if (iflag == 2) {
		fcn_resnet101 = new GeologicalHazardSeg(modelPath, 16, 1024, 1024);
	}

	SegmentContent content;

	cv::Mat img_bgr = cv::imread(inputFileName);
	if (img_bgr.empty())
	{
		return false;
	}
	fcn_resnet101->detect(img_bgr, content);

	if (content.flag)
	{		
		GdalOpencvTensor gdalOpencv;
		if (iflag == 1) {
			gdalOpencv.save2shp(inputFileName, content.class_mat, saveFileName,1);
		}
		if (iflag == 2) {
			gdalOpencv.save2shp(inputFileName, content.class_mat, saveFileName, 2);
		}

	
	}

	delete fcn_resnet101;
	return true;
}
bool DefectDetection::BuildingExtraction(string modelPath, string inputFileName, string saveFileName, string tmpPath) {

	BuildingExtractor extractor;
	extractor.initialize(modelPath); // 初始化模型

	// 影像是否存在
	QFile file(QString::fromStdString(inputFileName));
	if (!file.exists()) {
		return false;
	}

	//读取当前影像投影信息
	GDALAllRegister(); //注册所有的驱动;
	CPLSetConfigOption("GDAL_FILENAME_IS_UTF8", "YES");	// 支持中文路径;
	GDALDataset* poDataset;   //GDAL数据集;
	poDataset = (GDALDataset*)GDALOpen(inputFileName.c_str(), GA_ReadOnly); //读取影像;
	if (poDataset == NULL)//为空时表示读取失败;
	{
		cout << "fail in open files!!!" << endl;
		return false;
	}
	double GeoTransform[6];
	poDataset->GetGeoTransform(GeoTransform);
	int imgH = poDataset->GetRasterYSize();
	int imgW = poDataset->GetRasterXSize();
	int nBands = poDataset->GetRasterCount();
	GDALDataType dataType = poDataset->GetRasterBand(1)->GetRasterDataType();
	CPLSetConfigOption("GDAL_FILENAME_IS_UTF8", "NO");
	poDataset->Close();

	//if (!((nBands == 4 && dataType == GDT_UInt16) || (dataType == GDT_Byte && nBands == 3) || (dataType == GDT_UInt16 && nBands == 3))) {
	//	return false;
	//}
	// 判断是否需要裁图 
	if ( (imgW >= 1024 || imgH >= 1024 ) && 
		   !(  GeoTransform[0] == 0
			&& GeoTransform[1] == 1
			&& GeoTransform[2] == 0
			&& GeoTransform[3] == 0
			&& GeoTransform[4] == 0
			&& GeoTransform[5] == 1)) {
		// 开始裁图
		QString timestampStr = QString::number(QDateTime::currentDateTime().toSecsSinceEpoch()); // 获取当前时间戳
		QString saveCropPath = QString::fromStdString(tmpPath)+timestampStr; // 保存裁剪图像的路径

		QDir(saveCropPath).mkdir(saveCropPath);

		QString saveCropTifPath = saveCropPath + "/tif"; // 保存裁剪tif的路径
		QDir(saveCropTifPath).mkdir(saveCropTifPath);

		int CropSize = 1024;
		double RepetitionRate = 0.02;
		GdalOpencvTensor got;
		got.TifCrop(inputFileName, saveCropTifPath.toStdString(), CropSize, RepetitionRate);

		QString saveCropPngPath = saveCropPath + "/png"; // 保存裁剪png的路径
		QDir(saveCropPngPath).mkdir(saveCropPngPath);

		QString saveCropResultPath = saveCropPath + "/result"; // 保存裁剪result的路径
		QDir(saveCropResultPath).mkdir(saveCropResultPath);

		QString saveCropResultTifPath = saveCropPath + "/resultTif"; // 保存裁剪result tif的路径
		QDir(saveCropResultTifPath).mkdir(saveCropResultTifPath);

		QStringList tifPathList = GetImages(saveCropTifPath); // 转png
		for (QString tifpath : tifPathList) {
				QFileInfo tifinfo(tifpath);
				QString originpath = saveCropTifPath + "/" + tifpath; // 切片tif路径
				QString pngpath = saveCropPngPath + "/" + tifinfo.baseName() + ".png"; // 切片png路径
				QString resultpath = saveCropResultPath + "/" + tifinfo.baseName() + ".png"; // 切片result路径
				got.renderTiffExample((saveCropTifPath + "/" + tifpath).toStdString(), pngpath.toStdString());
				cv::Mat originPng = cv::imread(pngpath.toStdString());
				if (originPng.empty()) {
					return false;
				}
				// 开始提取
				cv::Mat cropresult = extractor.extract(originPng);
				// 保存结果 
				cv::imwrite(resultpath.toStdString(), cropresult);
			}

		got.arrtoras(saveCropTifPath.toStdString(), saveCropResultPath.toStdString(), saveCropResultTifPath.toStdString()); // 赋坐标
		got.mergeTiffImages(saveCropResultTifPath.toStdString(), saveCropResultTifPath.toStdString() + "_result.tif");
		bool re = got.save2shp(
			inputFileName ,
			saveCropResultTifPath.toStdString() + "_result.tif", 
			saveCropPath.toStdString(), 
			saveFileName,
			2,
			GeoTransform[1] * 2); // 这里只有smoothing
		return re;
	}

	
	if (imgW <= 1024 && imgH <= 1024 
		&& GeoTransform[0] == 0 
		&& GeoTransform[1] == 1
		&& GeoTransform[2] == 0
		&& GeoTransform[3] == 0
		&& GeoTransform[4] == 0
		&& GeoTransform[5] == 1) { // 无坐标 小图 0 1 0 0 0 1
		cv::Mat image = cv::imread(inputFileName);
		image = this->padImage(image, 1024); // padding 图片
		if (image.empty()) {
			return false;
		}
		cv::Mat result = extractor.extract(image);
		cv::imwrite(saveFileName + "_result.png", result); 
		GdalOpencvTensor gdalOpencv; 
		gdalOpencv.save2shp(inputFileName, result, saveFileName, 2, 0.01);
		return true;
	}

	return false; 
}
bool DefectDetection::SAMByShp(string tifPath, string shpPath, string savePath)
{
	std::cout << tifPath << std::endl;
	std::cout << shpPath << std::endl;
	std::cout << savePath << std::endl;
	// 影像是否存在
	QFile tifFile(QString::fromStdString(tifPath));
	if (!tifFile.exists()) {
		return false;
	}
	// shp是否存在
	QFile shpFile(QString::fromStdString(shpPath));
	if (!shpFile.exists()) {
		return false;
	}

	//读取当前影像投影信息
	GDALAllRegister(); //注册所有的驱动;
	CPLSetConfigOption("GDAL_FILENAME_IS_UTF8", "YES");	// 支持中文路径;
	GDALDataset* tifDataset = (GDALDataset*)GDALOpen(tifPath.data(), GA_ReadOnly); //读取影像;
	if (tifDataset == NULL)//为空时表示读取失败;
	{
		return false;
	}
	// 获取宽高
	int tifW = tifDataset->GetRasterXSize();
	int tifH = tifDataset->GetRasterYSize();

	// 读取坐标信息
	double geoTransform[6];
	CPLErr tifErr = tifDataset->GetGeoTransform(geoTransform);
	if (tifErr != CE_None) {
		tifDataset->Close();
		return false;
	}
	tifDataset->Close();

	GDALDataset* shpDataset = (GDALDataset*)GDALOpenEx(shpPath.data(), GDAL_OF_VECTOR, nullptr, nullptr, nullptr); //读取shp文件
	if (shpDataset == NULL || shpDataset->GetLayerCount() != 1) {
		shpDataset->Close();
		return false;
	}

	OGRLayer* layer = shpDataset->GetLayer(0);
	//for(int i=0;i<layer->GetRe)
	std::vector<PointInfo> points;
	std::vector<BoxInfo> boxes;
	OGRFeature* feature;
	while (feature = layer->GetNextFeature()) {
		 //= layer->GetNextFeature();
		if (feature == NULL)
			break;
		OGRGeometry* geometry = feature->GetGeometryRef();
		if (geometry == NULL) {
			delete geometry;
			continue;
		}
		OGRPoint point;
		OGRErr pointErr = geometry->Centroid(&point);
		if (pointErr != OGRERR_NONE) {
			break;
		}
		PointInfo pt;
		pt.positive = points.size() == 0 ;
		pt.pt.x = ((double)point.getX() - geoTransform[0]) / geoTransform[1];
		pt.pt.y = ((double)point.getY() - geoTransform[3]) / geoTransform[5];
		if (pt.pt.x < 0 || pt.pt.y < 0 || pt.pt.x > tifW || pt.pt.y > tifH) {
			continue;
		}
		//std::cout << pt.pt.x << " " << pt.pt.y << std::endl;
		points.push_back(pt);
	}

	// 读取图像信息
	cv::Mat tifMat = cv::imread(tifPath);
	// 开始提取
	this->cfg->HasMaskInput = false;
	this->samOnnxRunner->InferenceSingleImage(*(this->cfg), tifMat, points, boxes);
	this->cfg->HasMaskInput = true;
	cv::Mat mask = this->samOnnxRunner->InferenceSingleImage(*(this->cfg), tifMat, points, boxes)[0].mask;

	// 将其转换成shp
	GdalOpencvTensor gdalOpencv;
	std::string maskTifPath = savePath + "result.tif";

	std::string maskShpMidPath1 = this->cfg->SaveDir + "/ShpMid1.shp"; 
	std::string maskShpMidPath2 = this->cfg->SaveDir + "/ShpMid2.shp";
	std::string maskShpMidPath3 = savePath + "_ResultShp.shp";
	GDALDataset* maskTifDataset = gdalOpencv.Mat2File(maskTifPath, mask, geoTransform, 1); // 数组转栅格

	gdalOpencv.tif2shp(maskTifDataset,maskShpMidPath1, 2); // 栅格转shp
	gdalOpencv.smoothing(maskShpMidPath1.c_str(), maskShpMidPath2.c_str(), geoTransform[1]);
	gdalOpencv.simplify_geometry(maskShpMidPath2.c_str(), maskShpMidPath3.c_str(), geoTransform[1]);

	// 将其转换成shp
	//GdalOpencvTensor gdalopencv;

	delete layer;
	delete feature;

	return true;
}

bool DefectDetection::RoadRepair(string modelPath, string inputFileName, string maskFileName, string saveFileName, double buffer_distance) {
	try {
		// 初始化修复器
		RoadRepairer repairer;
		repairer.initialize(modelPath);

		// 读取输入图像
		cv::Mat image = cv::imread(inputFileName);
		if (image.empty()) {
			std::cerr << "Failed to read input image: " << inputFileName << std::endl;
			return false;
		}

		// 读取遮罩图像
		cv::Mat coarseMask = cv::imread(maskFileName, cv::IMREAD_GRAYSCALE);
		if (coarseMask.empty()) {
			std::cerr << "Failed to read mask image: " << maskFileName << std::endl;
			return false;
		}

		// 执行道路修复
		cv::Mat result = repairer.repair(image, coarseMask);

		// 保存结果图像
		cv::imwrite(saveFileName + "_result.png", result);

		// 调用GdalOpencvTensor保存为shp文件
		GdalOpencvTensor gdalOpencv;
		gdalOpencv.save2shp(inputFileName, result, saveFileName, 2, buffer_distance);

		return true;
	}
	catch (const std::exception& e) {
		std::cerr << "Error in RoadExtraction: " << e.what() << std::endl;
		return false;
	}
}

QStringList DefectDetection::GetImages(QString file_path)
{
	QStringList qImagesList;
	QDir folder(file_path);
	if (!folder.exists()) return qImagesList;
	QStringList filters;
	filters << "*.tif" << "*.tiff" << "*.jpg" << "*.png" << "*.img";
	folder.setNameFilters(filters);
	int imagesCount = folder.count();
	for (int i = 0; i < imagesCount; i++)
	{
		QString imageName = folder[i];
		qImagesList.push_back(imageName);
	}
	return qImagesList;
}

cv::Mat DefectDetection::padImage(const cv::Mat& input, int targetSize)
{
	// 获取输入图片的宽高
	int originalWidth = input.cols;
	int originalHeight = input.rows;

	// 如果图片尺寸已经大于等于目标尺寸，直接返回原图
	if (originalWidth >= targetSize && originalHeight >= targetSize) {
		return input.clone();
	}

	// 计算需要填充的宽度和高度
	int paddedWidth = max(originalWidth, targetSize);
	int paddedHeight = max(originalHeight, targetSize);

	// 创建填充后的图像（用0填充）
	cv::Mat paddedImage = cv::Mat::zeros(paddedHeight, paddedWidth, input.type());

	// 从左上角开始复制原图
	cv::Rect roi(0, 0, originalWidth, originalHeight);
	input.copyTo(paddedImage(roi));

	return paddedImage;
}

int DefectDetection::checkAndProcessImage(const string& imagePath)
{
	// 注册所有GDAL驱动
	GDALAllRegister();

	// 打开数据集
	GDALDataset* dataset = (GDALDataset*)GDALOpen(imagePath.c_str(), GA_ReadOnly);
	if (dataset == nullptr) {
		cout << "Error: Cannot open image file." << endl;
		return 0;
	}

	// 获取波段数
	int bandCount = dataset->GetRasterCount();

	// 获取第一个波段（用于判断数据类型）
	GDALRasterBand* band = dataset->GetRasterBand(1);
	if (band == nullptr) {
		cout << "Error: Cannot get raster band." << endl;
		GDALClose(dataset);
		return 0;
	}

	// 获取数据类型
	GDALDataType dataType = band->GetRasterDataType();

	// 检查是否有地理坐标
	bool hasGeoreference = false;
	double geoTransform[6];
	if (dataset->GetGeoTransform(geoTransform) == CE_None) {
		// 检查变换参数是否有效（不全为0）
		if (geoTransform[0] != 0.0 || geoTransform[1] != 1.0 ||
			geoTransform[2] != 0.0 || geoTransform[3] != 0.0 ||
			geoTransform[4] != 0.0 || geoTransform[5] != 1.0) {
			hasGeoreference = true;
		}
	}

	// 判断条件并执行相应处理
	if (hasGeoreference) {
		if (dataType == GDT_UInt16 && bandCount == 4) {
			return 2;
		}
		// 8位 3波段情况
		else if (dataType == GDT_Byte && bandCount == 3) {
			return 3;
		}
		else {
			cout << "Image format does not match required specifications." << endl;
			cout << "Current format: " << endl;
			cout << "Bit depth: " << GDALGetDataTypeName(dataType) << endl;
			cout << "Band count: " << bandCount << endl;
		}
	}
	else {
		cout << "Image does not contain georeferencing information." << endl;
		return 4;
	}

	// 清理资源
	GDALClose(dataset);
	return 1;
}
